package edu.cmu.cs211.compression.lzw;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Hashtable;

import edu.cmu.cs211.compression.Compressor;
import edu.cmu.cs211.compression.io.BitReader;
import edu.cmu.cs211.compression.io.BitWriter;
import edu.cmu.cs211.compression.lzw.LempelZivCodeReader;
import edu.cmu.cs211.compression.lzw.LempelZivCodeWriter;

/**
 * A compressor implementing LZW compression
 */
public class LempelZivCompressor extends Compressor {

	/**
	 * A tree node
	 * @author Mita Yun
	 * Each tree node contains a hashtable of treenodes,
	 * which can be connected to other treenodes
	 *
	 */
	public class LZWTreeNode
	{
		private Hashtable<Byte,LZWTreeNode> ht;
		private Byte data;
		private int index;	// code of data
		//start off with a dictionary of all the values..
		public LZWTreeNode()
		{
			ht = new Hashtable<Byte,LZWTreeNode>();
			for(int i = 0; i < 256; ++i)
			{
				ht.put(Byte.valueOf((byte)i), new LZWTreeNode(Byte.valueOf((byte)i), i));
			}
		}

		public LZWTreeNode(Byte data, int index)
		{
			ht = new Hashtable<Byte,LZWTreeNode>();
			this.data = data;
			this.index = index;
		}
		
		public LZWTreeNode getNodeByData(Byte data)
		{
			return ht.get(data);
		}
		
		/*
		 * make a new node with code index
		 * and insert it into the current level of hashtable
		 */
		public void insertNode(Byte data, int index)
		{
			LZWTreeNode newNode = new LZWTreeNode(data, index);
			if(ht.contains(data))
				throw new IllegalArgumentException();
			ht.put(data, newNode);
		}
		
		public void setData(Byte data)
		{
			this.data = data;
		}
		
		public Byte getData()
		{
			return this.data;
		}
		
		public void setIndex(int index)
		{
			this.index = index;
		}
		
		public int getIndex()
		{
			return this.index;
		}	
	}
	
	/** The number of bits in each LZW code */
	private static final int MAX_LZW_OUTPUT_WIDTH = 31;
	/** The maximum number of codes */
	private static final int MAX_LZW_CODES = (1 << MAX_LZW_OUTPUT_WIDTH) - 1;

	// @see Compressor#compress(io.BitReader, io.BitWriter)
	@Override
	public void compress(BitReader in, BitWriter out) throws IOException 
	{
		if(in == null || out == null)
			throw new NullPointerException();
		
		// code after the first level of tree
		int index = 256;
		
		//root is a hashtable
		LZWTreeNode root = new LZWTreeNode();
		
		int readBytes;
		
		LempelZivCodeWriter lzwWriter = new LempelZivCodeWriter(out,256);
		
		LZWTreeNode search = root;	// current node to look for
		LZWTreeNode previousNode = null;	// previous node that's searched
		
		
		while((readBytes = in.readByte()) != -1)
		{
			Byte readByte = Byte.valueOf((byte)readBytes);
		
			if((search = search.getNodeByData(readByte)) != null)
			{
				previousNode = search;
			}
			else
			{
				//we traversed to an emtpy spot in the tree to insert the new value..
				lzwWriter.writeCode(previousNode.getIndex());
				
				previousNode.insertNode(readByte, index++);
				
				if(index > MAX_LZW_CODES)
					throw new IllegalArgumentException();
			
				//since we already have the "next" byte that we should be processing, begin searching our tree
				search = root.getNodeByData(readByte);
				previousNode = search;
			}
		}
		// reached the end of the stream, write previous node
		if(previousNode != null)
		{
			lzwWriter.writeCode(previousNode.getIndex());
		}
		
		lzwWriter.flush();
	}

	// @see Compressor#expand(io.BitReader, io.BitWriter)
	@Override
	public void expand(BitReader in, BitWriter out) throws IOException {
		if(in == null || out == null)
			throw new NullPointerException();
		
		//make a "string" array to store our dictionary.. we can also use an arraylist of bytes..
		ArrayList<ArrayList<Byte>> LZWDictionary = new ArrayList<ArrayList<Byte>>();
		
		for(int i = 0; i < 256; ++i)
		{
			ArrayList<Byte> b = new ArrayList<Byte>();
			b.add((byte)i);
			LZWDictionary.add(b);
		}
		
		LempelZivCodeReader LZWReader = new LempelZivCodeReader(in, 256);
		
		int readCode;
		
		ArrayList<Byte> decodedWord = null;
		
		//read the first value..
		if((readCode = LZWReader.readCode()) != -1)
			decodedWord = LZWDictionary.get(readCode);
		else
			return;
		
		if(decodedWord == null)
		{
			return;
		}
		
		//print it out..
		for (int i = 0; i < decodedWord.size(); i++)
		{
		     out.writeByte(decodedWord.get(i));
		}
		
		while((readCode = LZWReader.readCode()) != -1)
		{
			// case 1, the next section that we are getting is NOT in the dictionary..
			if(readCode == LZWDictionary.size())
			{
				ArrayList<Byte> newArrayList = new ArrayList<Byte>();
				for (Byte b : decodedWord)
					newArrayList.add(b);
				newArrayList.add(decodedWord.get(0));

				LZWDictionary.add(newArrayList);
				decodedWord = LZWDictionary.get(readCode);
				for (int i = 0; i < decodedWord.size(); i++)
				{
		            out.writeByte(decodedWord.get(i));
				}
			}
			
			// case 2, the next section that we are getting is in the dictionary..
			else if (readCode < LZWDictionary.size())
			{
				ArrayList<Byte> newArrayList = new ArrayList<Byte>();
				for (Byte b : decodedWord)
					newArrayList.add(b);
				newArrayList.add(LZWDictionary.get(readCode).get(0));
				
				LZWDictionary.add(newArrayList);
				decodedWord = LZWDictionary.get(readCode);
				for (int i = 0; i < decodedWord.size(); i++)
				{
		            out.writeByte(decodedWord.get(i));
				}
			}
			else
				throw new IOException();
		}
		out.flush();
	}
}
